/**
 * 
 */
package edu.umd.clip.lm.factors;

import java.util.*;
import org.w3c.dom.*;
import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public final class FactorTupleDescription extends XMLizable {
	public final static String XML_ELEMENT_NAME = "factor-tuple-description";
	public final static String XML_ATTR_MAIN_FACTOR_NAME = "main-factor";
	
	private FactorDescription[] descriptions = null;
	private HashMap<String, Byte> factorid2idx = null;
//	private boolean locked = false;
	private byte[] hiddenFactors = null;
	private byte[] overtFactors = null;
	private byte[] parentIndex = null;
	private byte[] factorBitStart = null;
	private byte[] factorBitLength = null;
	private byte compactHiddenFactorsSize = 0; // bits
	private byte compactOvertFactorsSize = 0; // bits
	private byte mainFactorIndex = -1;
	private String mainFactor = "W";
	private Map<FactorTuple,FactorTuple> allOvertFactors = null;
	private long overtFactorsMask;
	private long hiddenFactorsMask;
	
	public FactorTupleDescription(FactorDescription[] descriptions) {
		init(descriptions);
	}
	
	private void init(FactorDescription[] descriptions) {
		assert(descriptions.length <= Byte.MAX_VALUE);
		
		this.descriptions = descriptions;
		factorid2idx = new HashMap<String, Byte>(descriptions.length);
		
		ArrayList<Byte> hiddens = new ArrayList<Byte>(0);
		ArrayList<Byte> overts = new ArrayList<Byte>(0);
		
		factorBitStart = new byte[this.descriptions.length];
		factorBitLength = new byte[this.descriptions.length];

		byte bitOffset = 0;
		for(byte i = 0; i<this.descriptions.length; ++i) {
			FactorDescription desc = this.descriptions[i];
			Byte idx = new Byte(i);
			if (desc.getBits() == 0) {
				desc.setBits(desc.computeBitSize());
			}
			long mask = (1L << desc.getBits()) - 1;
			mask <<= bitOffset;
			
			if (desc.isOvert()) {
				overts.add(idx);
				overtFactorsMask |= mask;
				compactOvertFactorsSize += desc.getBits();
			} else {
				hiddens.add(idx);
				hiddenFactorsMask |= mask;
				compactHiddenFactorsSize += desc.getBits();
			}
			factorid2idx.put(desc.getId(), idx);
			
			factorBitStart[i] = bitOffset;
			factorBitLength[i] = desc.getBits();
			bitOffset += desc.getBits();
			assert(bitOffset <= Long.SIZE);
		}
		assert(overts.size() > 0);
		mainFactorIndex = factorid2idx.get(mainFactor);
		
		hiddenFactors = new byte[hiddens.size()];
		for(byte i = 0; i<hiddens.size(); ++i) {
			hiddenFactors[i] = hiddens.get(i).byteValue();
		}
		overtFactors = new byte[overts.size()];
		for(byte i = 0; i<overts.size(); ++i) {
			overtFactors[i] = overts.get(i).byteValue();
		}
		
		parentIndex = new byte[this.descriptions.length];
		for(int i=0; i<parentIndex.length; ++i) {
			if (this.descriptions[i].getParent() == null) {
				parentIndex[i] = -1;
			} else {
				for(byte j=0; j<parentIndex.length; ++j) {
					if (this.descriptions[j] == this.descriptions[i].getParent()) {
						parentIndex[i] = j;
						break;
					}
				}
			}
		}
		
		FactorTuple.overtMask = overtFactorsMask;
		FactorTuple.hiddenMask = hiddenFactorsMask;
	}
	
	static public FactorTupleDescription fromXML(Element xmlnode) {
		return new FactorTupleDescription(xmlnode);
	}
	
	public FactorTupleDescription(Element xmlnode) {
		HashMap<String, FactorDescription> descriptionsById = new HashMap<String, FactorDescription>(10);
		ArrayList<FactorDescription> descriptions = new ArrayList<FactorDescription>(10);
		
		NodeList factorDescList = XMLUtils.getChildrenElementsByTagName(xmlnode, FactorDescription.XML_ELEMENT_NAME);
		for(int i = 0; i < factorDescList.getLength(); ++i) {
			Element node = (Element)factorDescList.item(i);
			if (node.getNodeType() != Node.ELEMENT_NODE) continue;
			NamedNodeMap attrs = node.getAttributes();
			String id = attrs.getNamedItem(FactorDescription.XML_ATTR_ID_NAME).getNodeValue();
			String name = attrs.getNamedItem(FactorDescription.XML_ATTR_NAME_NAME).getNodeValue();
			boolean overt = attrs.getNamedItem(FactorDescription.XML_ATTR_IS_OVERT_NAME).getNodeValue().equals("yes");
			
			FactorDescription parentDesc = null;
			Node parentNode = attrs.getNamedItem(FactorDescription.XML_ATTR_PARENT_ID_NAME);
			if (parentNode != null) {
				String parentId = parentNode.getNodeValue();
				parentDesc = descriptionsById.get(parentId);
				if (parentDesc == null) {
					throw new Error("Undefined Factor Description '"+parentId+"' is referenced in '"+id+"'.");
				}
			}
			FactorDescription fd = new FactorDescription(id, name, parentDesc);
			fd.setOvert(overt);
			fd.fromXML(node);
			
			descriptionsById.put(id, fd);
			descriptions.add(fd);
		}
		
		String mainFactor = xmlnode.getAttribute(XML_ATTR_MAIN_FACTOR_NAME);
		if (!mainFactor.equals("")) {
			if (descriptionsById.containsKey(mainFactor)) {
				this.mainFactor = mainFactor;
			} else {
				throw new Error("The main factor '"+mainFactor+"' is not found.");
			}
		}
		init(descriptions.toArray(new FactorDescription[0]));
	}
	
	public Element createXML(Document doc) {
		Element elem = doc.createElement(XML_ELEMENT_NAME);
		for(int i=0; i<descriptions.length; ++i) {
			elem.appendChild(descriptions[i].createXML(doc));
		}
		return elem;
	}
	
	public long createTuple() {
		return 0;
	}
	
	public long createStartTuple() {
		int[] values = new int[numFactors()];
		Arrays.fill(values, Dictionary.getStart());
		return createTuple(values);
	}
	
	public long createEndTuple() {
		int[] values = new int[numFactors()];
		Arrays.fill(values, Dictionary.getEnd());
		return createTuple(values);
	}
	
	public long createTuple(int[] values) {
		if (values == null) {
			values = new int[numFactors()];
			Arrays.fill(values, Dictionary.getNull());
		} else {
			assert(values.length == numFactors());
		}
		long tuple = 0;
		for(byte i=0; i<values.length; ++i) {
			tuple = FactorTuple.setValue(tuple, i, values[i]);
		}
		return tuple;
	}
	
	public FactorDescription[] getDescriptions() {
		return descriptions;
	}

	public byte[] getHiddenFactors() {
		return hiddenFactors;
	}

	public byte[] getOvertFactors() {
		return overtFactors;
	}
	
	public long getOvertFactorsMask() {
		return overtFactorsMask;
	}
	
	public long getHiddenFactorsMask() {
		return hiddenFactorsMask;
	}
	
	public FactorDescription getDescription(byte idx) {
		return descriptions[idx];
	}
	
	public Dictionary getDictionary(byte idx) {
		return descriptions[idx].getDictionary();
	}
	
	public Dictionary getDictionary(String factorId) {
		return descriptions[getFactorIndex(factorId)].getDictionary();
	}
	
	public byte numFactors() {
		return (byte) descriptions.length;
	}
	
	/**
	 * @param id the id of the factor
	 * @return index of the factor or -1 if there is no such factor
	 */
	public byte getFactorIndex(String id) {
		Byte idx = factorid2idx.get(id);
		if (idx == null) return -1;
		return idx.byteValue();
	}
	
	public byte getParentIndex(byte idx) {
		return parentIndex[idx];
	}

	public String printOvertFactors(FactorTuple factors) {
		return printOvertFactors(factors.getBits());
	}
	
	public String printOvertFactors(long bits) {
		FactorTuple tuple = new FactorTuple(bits & overtFactorsMask);
		return tuple.toStringNoNull();
	}
	
	public String printHiddenFactors(FactorTuple factors) {
		return printHiddenFactors(factors.getBits());
	}
	
	public String printHiddenFactors(long bits) {
		FactorTuple tuple = new FactorTuple(bits & hiddenFactorsMask);
		return tuple.toStringNoNull();
	}
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.util.XMLizable#getElementName()
	 */
	@Override
	public String getElementName() {
		return XML_ELEMENT_NAME;
	}
	
	public final Map<FactorTuple,FactorTuple> getAllOvertFactors() {
		if (allOvertFactors != null) return allOvertFactors;
		synchronized(this) {
			if (allOvertFactors == null) {
				allOvertFactors = buildAllOvertFactors();
			}
		}
		return allOvertFactors;
	}
	
	private Map<FactorTuple,FactorTuple> buildAllOvertFactors() {
		Dictionary mainDict = getDictionary(mainFactorIndex);
		HashMap<FactorTuple,FactorTuple> factors = new HashMap<FactorTuple,FactorTuple>(mainDict.size());
		
		for(DictionaryIterator it = mainDict.iterator(true); it.hasNext(); ) {
			FactorTuple tuple = new FactorTuple();
			tuple.setValue(mainFactorIndex, it.next());
			
			if (overtFactors.length == 1) {
				factors.put(tuple,tuple);
				continue;
			}
			
			buildAllOvertFactorsHelper(factors, tuple, (byte)0);
		}
		return factors;
	}
	
	private void buildAllOvertFactorsHelper(HashMap<FactorTuple,FactorTuple> factors, FactorTuple tuple, byte pos) {
		if (pos == overtFactors.length) {
			FactorTuple f = new FactorTuple(tuple.getBits());
			factors.put(f, f);
			return;
		}
		if (overtFactors[pos] == mainFactorIndex) {
			buildAllOvertFactorsHelper(factors, tuple, ++pos);
			return;
		}
		
		byte idx = overtFactors[pos];
		int[] vals = descriptions[idx].dependentOvertFactors(tuple);
		if (vals == null) {
			buildAllOvertFactorsHelper(factors ,tuple, ++pos);
			return;
		}

		++pos;
		for(int v : vals) {
			tuple.setValue(idx, v);
			buildAllOvertFactorsHelper(factors ,tuple, pos);
		}
	}
	
	public byte getMainFactorIndex() {
		return mainFactorIndex;
	}

	public String getMainFactor() {
		return mainFactor;
	}

	public final byte[] getFactorBitStart() {
		return factorBitStart;
	}

	public final byte[] getFactorBitLength() {
		return factorBitLength;
	}
	
	public byte getCompactHiddenFactorsSize() {
		return compactHiddenFactorsSize;
	}
	
	public byte getCompactOvertFactorsSize() {
		return compactOvertFactorsSize;
	}
	
	public byte getCompactHiddenFactorsBytes() {
		return (byte) ((compactHiddenFactorsSize + 7) >> 3);
	}
	
	public byte getCompactOvertFactorsBytes() {
		return (byte) ((compactOvertFactorsSize + 7) >> 3);
	}
	
	public byte[] packHiddenFactors(long factors) {
		byte bytes[] = new byte[getCompactHiddenFactorsBytes()];
		byte factorOffset = 0;
		for(byte factor : hiddenFactors) {
			long bits = factors >> factorBitStart[factor];
			for(byte i=0; i<factorBitLength[factor]; ++i) {
				if ((bits & (1 << i)) != 0) {
					int offset = (i + factorOffset) / 8;
					int bit = (i + factorOffset) % 8;
					bytes[offset] |= 1 << bit;
				}
			}
			factorOffset += factorBitLength[factor];
		}
		return bytes; 
	}
	
	public int packHiddenFactorsToInt(long factors) {
		int packed = 0;
		byte factorOffset = 0;
		for(byte factor : hiddenFactors) {
			long bits = factors >> factorBitStart[factor];
			for(byte i=0; i<factorBitLength[factor]; ++i) {
				if ((bits & (1 << i)) != 0) {
					// TODO: could be faster
					packed |= 1 << (factorOffset + i);
				}
			}
			factorOffset += factorBitLength[factor];
		}
		return packed; 
	}
	
	public long unpackHiddenFactors(byte[] bytes) {
		long factors = 0;
		byte factorOffset = 0;

		for(byte factor : hiddenFactors) {
			long bits = 0;
			for(short i=0; i<factorBitLength[factor]; ++i) {
				int offset = (i + factorOffset) / 8;
				int bit = (i + factorOffset) % 8;
				if ((bytes[offset] & (1 << bit)) != 0) {
					bits |= 1 << i;
				}
			}
			factors |= bits << factorBitStart[factor];
			factorOffset += factorBitLength[factor];
		}
		
		return factors;
	}
	
	public long unpackHiddenFactorsFromInt(int packed) {
		long factors = 0;
		byte factorOffset = 0;

		for(byte factor : hiddenFactors) {
			long bits = 0;
			for(short i=0; i<factorBitLength[factor]; ++i) {
				// TODO: could be faster
				if ((packed & (1 << (factorOffset + i))) != 0) {
					bits |= 1 << i;
				}
			}
			factors |= bits << factorBitStart[factor];
			factorOffset += factorBitLength[factor];
		}
		
		return factors;
	}

	public byte[] packOvertFactors(long factors) {
		byte bytes[] = new byte[getCompactOvertFactorsBytes()];
		byte factorOffset = 0;
		for(byte factor : overtFactors) {
			long bits = factors >> factorBitStart[factor];
			for(byte i=0; i<factorBitLength[factor]; ++i) {
				if ((bits & (1 << i)) != 0) {
					int offset = (i + factorOffset) / 8;
					int bit = (i + factorOffset) % 8;
					bytes[offset] |= 1 << bit;
				}
			}
			factorOffset += factorBitLength[factor];
		}
		return bytes; 
	}
	
	public int packOvertFactorsToInt(long factors) {
		int packed = 0;
		byte factorOffset = 0;
		for(byte factor : overtFactors) {
			long bits = factors >> factorBitStart[factor];
			for(byte i=0; i<factorBitLength[factor]; ++i) {
				if ((bits & (1 << i)) != 0) {
					// TODO: could be faster
					packed |= 1 << (factorOffset + i);
				}
			}
			factorOffset += factorBitLength[factor];
		}
		return packed; 
	}
	
	public long unpackOvertFactors(byte[] bytes) {
		long factors = 0;
		byte factorOffset = 0;

		for(byte factor : overtFactors) {
			long bits = 0;
			for(short i=0; i<factorBitLength[factor]; ++i) {
				int offset = (i + factorOffset) / 8;
				int bit = (i + factorOffset) % 8;
				if ((bytes[offset] & (1 << bit)) != 0) {
					bits |= 1 << i;
				}
			}
			factors |= bits << factorBitStart[factor];
			factorOffset += factorBitLength[factor];
		}
		
		return factors;
	}

	public long unpackOvertFactorsFromInt(int packed) {
		long factors = 0;
		byte factorOffset = 0;

		for(byte factor : overtFactors) {
			long bits = 0;
			for(short i=0; i<factorBitLength[factor]; ++i) {
				// TODO: could be faster
				if ((packed & (1 << (factorOffset + i))) != 0) {
					bits |= 1 << i;
				}
			}
			factors |= bits << factorBitStart[factor];
			factorOffset += factorBitLength[factor];
		}
		
		return factors;
	}

}
